#include <stdio.h>
#include <omp.h>

#include "mpi.h"
#include "oblivious.h"
#include "distributed_mult.h"
#include "stdlib.h"


// If this is false, the output is a CSV row.
static const bool VERBOSE = false;


int main(int argc, char *argv[]) {

  int a_rows;
  int a_cols;
  int b_cols;
  int base_case_size;

  if (argc < 5) {
    fprintf(stderr, "Wrong number of arguments (%d)\n", argc);
    fprintf(stderr, "usage: scaling a_rows a_cols b_cols base_case_size\n");
    fprintf(stderr, "Using default.\n");

    a_rows = 10;
    a_cols = 11;
    b_cols = 7;
    base_case_size = 2;

  } else {
    a_rows = atoi(argv[1]);
    a_cols = atoi(argv[2]);
    b_cols = atoi(argv[3]);
    base_case_size = atoi(argv[4]);
  }

  if (VERBOSE)
    printf("main: a_rows: %d; a_cols: %d; b_cols: %d; base_case_size: %d\n",
	   a_rows, a_cols, b_cols, base_case_size);
  
  MPI_Init(&argc, &argv);

  int my_rank, num_procs;
  MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
  MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

  double *my_a;
  double *my_b;
  double *my_c;

  /*
  int num_row_procs = (int)ceil(sqrt((double)num_procs));
  int num_col_procs = num_procs / num_row_procs;
  
  printf("row_procs: %d, col_procs: %d\n", num_row_procs, num_col_procs);
  */

  int my_a_rows, my_a_cols, my_b_rows, my_b_cols;
    
  my_a_rows = get_row_count(a_rows, my_rank, num_procs);
  my_a_cols = a_cols;
  my_b_rows = a_cols;
  my_b_cols = get_row_count(b_cols, my_rank, num_procs);
  
  my_a = (double*) malloc(sizeof(double) * my_a_rows * my_a_cols);
  my_b = (double*) malloc(sizeof(double) * my_b_rows * my_b_cols);
  my_c = (double*) malloc(sizeof(double) * a_rows * my_b_cols);

  //printf("my_a: %p, my_b: %p, my_c: %p\n", my_a, my_b, my_c);
  
  //printf("Allocated memory\n");

  for (int i = 0; i < my_a_rows * my_a_cols; i++) {
    
    my_a[i] = 1.0;
    
  }
  for (int i = 0; i < my_b_rows * my_b_cols; i++) {
    my_b[i] = 1.0;
  }

  for (int i = 0; i < a_rows * my_b_cols; i++) {
    my_c[i] = 0.0;
  }

  //printf("Finished writing memory\n");

  // start timing

  int my_c_start = 0;
  for (int i = 0; i < my_rank; i++) {
    my_c_start += get_row_count(a_rows, i, num_procs);
  }

  //printf("Starting multiplication\n");

  double start_time = MPI_Wtime();
  distributed_mat_mult_rec(my_a, my_b, my_c, a_rows, a_cols, b_cols, my_c_start, 
			   base_case_size);
  double end_time = MPI_Wtime();
  
  // end timing

  if (VERBOSE) {
    printf("Total time: %g\n", end_time - start_time);
    printf("my_c[0]: %g, proc: %d\n", my_c[0], my_rank);
    Matrix c_mat(my_c, a_rows, my_b_cols);
    c_mat.dump();
  }

  if (my_rank == 0) {
    printf("result: %d,%d,%d,%d,%d,%d,%f\n", num_procs, omp_get_max_threads(), 
	   a_rows, a_cols, b_cols, base_case_size, end_time - start_time);
  }

  MPI_Finalize();
    
  return 0;

}
